import numpy as np
import pickle
import torch

import gymnasium as gym


def evaluate_model(env_name, agent_name, model):
    # net = torch.load(model, weights_only=False)
    with open(model, 'rb') as f:
        net, state_norm = pickle.load(f)
    env = gym.make(env_name, render_mode='human')
    done = False
    state, _ = env.reset(seed=10)
    if state_norm is not None:
        state = state_norm(state, update=False)
    while not done:
        state = torch.tensor(state).unsqueeze(0).float()
        if agent_name in ['DQN']:
            q_values = net(state)
            action = torch.argmax(q_values).item()
        elif agent_name in ['DDPG', 'TD3', 'PPOContinuous']:
            action = net(state).detach().numpy().flatten()
        elif agent_name in ['PPODiscrete']:
            prob = net(state).detach().numpy().flatten()
            action = np.argmax(prob)
        else:
            raise ValueError(f"{agent_name} error")
        # action = env.action_space.sample()
        state, _, terminated, truncated, _ = env.step(action)
        if state_norm is not None:
            state = state_norm(state, update=False)
        done = terminated or truncated


if __name__ == "__main__":
    env_name = 'Humanoid-v5'
    agent_name = 'PPOContinuous'
    # evaluate_model(env_name, agent_name, f"../results/model/{env_name}.pth")
    evaluate_model(env_name, agent_name, f"../results/model/{env_name}_{agent_name}.pth")